### solves model 

initial_guess<-rep(3.0, n_students) ## initial guess for study time game

gamma_s_not_i<-function(s_not_i) (s_not_i)^(tau_s_not_i)
gamma_mu<-function(mu) exp(tau_mu_vec[1]*mu + tau_mu_vec[2]*mu^2)
  
#model test score
y<-function(s_i_avg,s_not_i_avg, mu_i) 
	{
	beta_vec[1] + beta_vec[2]*s_i_avg + beta_vec[3]*s_not_i_avg + mu_i
	}

calc_s_not_own<-function(A_array_this_month) function(s_vec)
	{
  n_students_tmp<-dim(A_array_this_month)[1]
  num_friends<-A_array_this_month %*% rep(1, n_students_tmp)
  tmp<-t(A_array_this_month %*% s_vec /num_friends )
  tmp<-as.numeric(tmp)
  tmp[num_friends==0]<-0
  
  tmp
	}

## deals with NA in observed study time
## could use this instead of calc_s_not_own, but this is slightly slower...
calc_s_not_own_data<-function(A_array_this_month) function(s_vec)
{
  s_NA_vec<- as.numeric(is.na(s_vec))  
  num_s_NA<-A_array_this_month %*% s_NA_vec ### count up NAs per student
  s_vec[is.na(s_vec)]<-0
  n_students_tmp<-dim(A_array_this_month)[1]
  num_friends<-A_array_this_month %*% rep(1, n_students_tmp)
  num_friends_no_s_NA<-num_friends - num_s_NA
  tmp<-t(A_array_this_month %*% s_vec /num_friends_no_s_NA )
  tmp<-as.numeric(tmp)
  tmp[num_friends==0]<-0
  
  tmp
}

# own effort (best response function)
## reduced-form psi, can map back to cost function
if (use_kappa_instead_of_theta==T){
  psi_to_vec<-function(mu,s_not_i) 
  {
    max(kappa_vec[1] + kappa_vec[2]/gamma_mu(mu) + kappa_vec[3]*gamma_s_not_i(s_not_i) + kappa_vec[4]*gamma_s_not_i(s_not_i)/gamma_mu(mu),0, na.rm=T)
  }
} else{
    if (use_conformity_spec==TRUE){
      psi_to_vec<-function(mu,s_not_i) 
      {
        max((beta_vec[2] - chi_vec[1])/(1 + chi_vec[4]) + (-chi_vec[2]/(1 + chi_vec[4]))/gamma_mu(mu) + (chi_vec[4]/(1 + chi_vec[4])) *gamma_s_not_i(s_not_i) + (chi_vec[4]*chi_vec[5]/(1 + chi_vec[4])) *gamma_s_not_i(s_not_i)/gamma_mu(mu) ,0, na.rm=T)
      }
    } else { # cost-reduction specification
      psi_to_vec<-function(mu,s_not_i) 
      {
      max(-theta_vec[3] - theta_vec[4]/gamma_mu(mu) + (beta_vec[2] - theta_vec[1])*gamma_s_not_i(s_not_i) - theta_vec[2]*gamma_s_not_i(s_not_i)/gamma_mu(mu),0, na.rm=T)
      }
    }
}
  
psi<-function(mu) function(s_not_i) 
  {
  Vectorize(psi_to_vec)(mu=mu, s_not_i=s_not_i)
}

# takes in theta_vec and iterates to solve for equilibrium
calc_equilibrium<-function(A_array_this_month) function(mu_vec) 
{
  iteration<-0
  dist<-1.0
  old_guess<-initial_guess[1:dim(A_array_this_month)[1]]
  while(dist>tol){
    new_guess<-psi(mu_vec)(calc_s_not_own(A_array_this_month)(old_guess))
    dist<- mean((new_guess - old_guess)^2)
    old_guess<-new_guess
    iteration<-iteration+1
  }
  list(ans=new_guess, iterations=iteration)
}

# takes in a vector of study times as initial guess and a person to increase by 1
calc_equilibrium_iterations<-function(A_array_this_month) function(mu_vec) function(initial_equilibrium) function(shocked_person_index, shock_delta)
{
  iteration<-0
  dist<-1.0
  shocked_person_perturbed<-initial_equilibrium[shocked_person_index] + shock_delta
  perturbation<- rep(0, length(initial_equilibrium))
  perturbation[shocked_person_index]<-shock_delta
  old_guess<-initial_equilibrium + perturbation
  ## array to collect s_vec per iteration
  s_vec_per_iteration<-initial_equilibrium
  while(dist>tol){
    new_guess<-psi(mu_vec)(calc_s_not_own(A_array_this_month)(old_guess))
    new_guess[shocked_person_index]<-shocked_person_perturbed
    dist<- mean((new_guess - old_guess)^2)
    old_guess<-new_guess
    iteration<-iteration+1
    s_vec_per_iteration<-cbind(s_vec_per_iteration, old_guess)
  }
  list(s_vec_per_iteration=s_vec_per_iteration)
}

## compute comparative statics given perturbed graph
compute_shock_comparative_statics<-function(A_array_this_month) function(mu_vec_y, mu_vec_study) function(initial_equilibrium) function(shocked_person_index, shock_delta)
{
  ## compute new equilibrium
  s_array_iterations<-calc_equilibrium_iterations(A_array_this_month)(mu_vec_study)(initial_equilibrium)(shocked_person_index=shocked_person_index, shock_delta=shock_delta)[[1]]
  
  ##first, last iterations
  s_vec_shocked<-s_array_iterations[,1]
  s_vec_first_iteration<-s_array_iterations[,2]
  s_vec_last_iteration<-s_array_iterations[,dim(s_array_iterations)[2]]
  s_not_i_vec_shocked<-calc_s_not_own(A_array_this_month)(s_vec_shocked)
  s_not_i_vec_first_iteration<-calc_s_not_own(A_array_this_month)(s_vec_first_iteration)
  s_not_i_vec_last_iteration<-calc_s_not_own(A_array_this_month)(s_vec_last_iteration)
  y_vec_shocked<-y(s_vec_shocked, s_not_i_vec_shocked, mu_vec_y)
  y_vec_first_iteration<-y(s_vec_first_iteration, s_not_i_vec_first_iteration, mu_vec_y)
  y_vec_last_iteration<-y(s_vec_last_iteration, s_not_i_vec_last_iteration, mu_vec_y)
  
  ## affected nodes
  neighborhood_order<-1
  A_graph_this_month<-graph.adjacency(A_array_this_month)
  V(A_graph_this_month)$names<-1:n_students
  connected_nodes<-neighborhood(A_graph_this_month, neighborhood_order, shocked_person_index)[[1]]
  connected_nodes_graph<-graph.neighborhood(A_graph_this_month, neighborhood_order, shocked_person_index)[[1]]
  
  # distances from shocked person to everyone in the network
  dist_from_shocked_person<-suppressWarnings(sapply(get.shortest.paths(A_graph_this_month, from=shocked_person_index)$vpath, FUN=function(x) length(x)-1))
  print(summary(dist_from_shocked_person))
  
  ## total produced, avg produced, other stats...
  ## plots showing graph, color shocked person, color effects by person
  print("GPA")
  print(summary( ((y_vec_first_iteration - y_vec_shocked)/shock_delta)[connected_nodes[-1]]))
  print(summary( ((y_vec_last_iteration -  y_vec_shocked)/shock_delta)[connected_nodes[-1]]))
  print("own study")
  print(summary( ((s_vec_first_iteration - s_vec_shocked)/shock_delta)[connected_nodes[-1]]))
  print(summary( ((s_vec_last_iteration -  s_vec_shocked)/shock_delta)[connected_nodes[-1]]))
  print("friend study")
  print(summary( ((s_not_i_vec_first_iteration - s_not_i_vec_shocked)/shock_delta)[connected_nodes[-1]]))
  print(summary( ((s_not_i_vec_last_iteration -  s_not_i_vec_shocked)/shock_delta)[connected_nodes[-1]]))
  print("total increases")
  print(c(y=sum((y_vec_last_iteration - y_vec_shocked)[connected_nodes[-1]]), s=sum((s_vec_last_iteration - s_vec_shocked)[connected_nodes[-1]]), s_not_i=sum((s_not_i_vec_last_iteration - s_not_i_vec_shocked)[connected_nodes[-1]])))
  
  ## summarize effects on s, s_not_i, y for students by neighborhood size
  max_neighborhood_order<-4
  summary_df<-data.frame(neighborhood_size=NA, iteration=NA, outcome=NA, value=NA)
    for (neighborhood_order in 1:max_neighborhood_order)
    {
      connected_nodes<-neighborhood(A_graph_this_month, neighborhood_order, shocked_person_index)[[1]]
      connected_nodes<-connected_nodes[-1]
      summary_df<-rbind(summary_df, c(neighborhood_order,"partial equil.","y", mean(((y_vec_first_iteration - y_vec_shocked)/shock_delta)[connected_nodes]))) 
      summary_df<-rbind(summary_df, c(neighborhood_order,"new equil.","y", mean((( y_vec_last_iteration - y_vec_shocked)/shock_delta)[connected_nodes]))) 
      
      summary_df<-rbind(summary_df, c(neighborhood_order,"partial equil.","s", mean(((s_vec_first_iteration - s_vec_shocked)/shock_delta)[connected_nodes]))) 
      summary_df<-rbind(summary_df, c(neighborhood_order,"new equil.","s", mean((( s_vec_last_iteration - s_vec_shocked)/shock_delta)[connected_nodes]))) 
      
      summary_df<-rbind(summary_df, c(neighborhood_order,"partial equil.","s_not_i", mean(((s_not_i_vec_first_iteration - s_not_i_vec_shocked)/shock_delta)[connected_nodes]))) 
      summary_df<-rbind(summary_df, c(neighborhood_order,"new equil.","s_not_i", mean((( s_not_i_vec_last_iteration - s_not_i_vec_shocked)/shock_delta)[connected_nodes]))) 
    }
  
  summary_df<-summary_df[-1,]
  summary_df$outcome<-factor(summary_df$outcome, as.character(summary_df$outcome))
  summary_df$iteration<-factor(summary_df$iteration, as.character(summary_df$iteration))
  summary_df$value<-as.numeric(summary_df$value)
  print(head(summary_df))
  
  y_range<-range(summary_df$value)

  ## gain from equil df, for changes in y
  equil_gain_df<-data.frame(neighborhood_size=NA, outcome=NA, pct_gain_over_partial=NA)
  
  for (neighborhood_order in 1:max_neighborhood_order)
  {
    equil_gain_df<-rbind(equil_gain_df, c(neighborhood_order, outcome="y", with(summary_df, summary_df[neighborhood_size==neighborhood_order & outcome=="y" & iteration=="new equil.",]$value / summary_df[neighborhood_size==neighborhood_order & outcome=="y" & iteration=="partial equil.",]$value)-1 ))
    
    equil_gain_df<-rbind(equil_gain_df, c(neighborhood_order, outcome="s", with(summary_df, summary_df[neighborhood_size==neighborhood_order & outcome=="s" & iteration=="new equil.",]$value / summary_df[neighborhood_size==neighborhood_order & outcome=="s" & iteration=="partial equil.",]$value)-1 ))
    
    equil_gain_df<-rbind(equil_gain_df, c(neighborhood_order, outcome="s_not_i", with(summary_df, summary_df[neighborhood_size==neighborhood_order & outcome=="s_not_i" & iteration=="new equil.",]$value / summary_df[neighborhood_size==neighborhood_order & outcome=="s_not_i" & iteration=="partial equil.",]$value)-1 ))
  }
  
  equil_gain_df<-equil_gain_df[-1,]
  equil_gain_df$outcome<-factor(equil_gain_df$outcome, as.character(equil_gain_df$outcome))
  equil_gain_df$pct_gain_over_partial<-as.numeric(equil_gain_df$pct_gain_over_partial)
  print(equil_gain_df)
}

## compute comparative statics given perturbed graph
## returns a dataframe with change in y, s, s_not_i, distance from shocked
compute_shock_comparative_statics_df<-function(A_array_this_month) function(mu_vec_y, mu_vec_study) function(initial_equilibrium) function(shocked_person_index, shock_delta)
{
  ## compute new equilibrium
  s_array_iterations<-calc_equilibrium_iterations(A_array_this_month)(mu_vec_study)(initial_equilibrium)(shocked_person_index=shocked_person_index, shock_delta=shock_delta)[[1]]
  
  ##first, last iterations
  s_vec_shocked<-s_array_iterations[,1]
  s_vec_first_iteration<-s_array_iterations[,2]
  s_vec_last_iteration<-s_array_iterations[,dim(s_array_iterations)[2]]
  s_not_i_vec_shocked<-calc_s_not_own(A_array_this_month)(s_vec_shocked)
  s_not_i_vec_first_iteration<-calc_s_not_own(A_array_this_month)(s_vec_first_iteration)
  s_not_i_vec_last_iteration<-calc_s_not_own(A_array_this_month)(s_vec_last_iteration)
  y_vec_shocked<-y(s_vec_shocked, s_not_i_vec_shocked, mu_vec_y)
  y_vec_first_iteration<-y(s_vec_first_iteration, s_not_i_vec_first_iteration, mu_vec_y)
  y_vec_last_iteration<-y(s_vec_last_iteration, s_not_i_vec_last_iteration, mu_vec_y)
  
  # distances from shocked person to everyone in the ENTIRE network
  A_graph_this_month<-graph.adjacency(A_array_this_month)
  V(A_graph_this_month)$names<-1:n_students
  dist_from_shocked_person<-sapply(get.shortest.paths(A_graph_this_month, from=shocked_person_index)$vpath, function(x) length(x)-1)
  print(summary(dist_from_shocked_person))
  
  ## initial gain df
  initial_gain_df<-data.frame(pid=student_data$pid, y_diff = y_vec_first_iteration - y_vec_shocked, s_diff = s_vec_first_iteration - s_vec_shocked, s_not_i_diff = s_not_i_vec_first_iteration - s_not_i_vec_shocked, dist_from_shocked=dist_from_shocked_person, iteration = "partial")
  final_gain_df<-data.frame(pid=student_data$pid, y_diff = y_vec_last_iteration - y_vec_shocked, s_diff = s_vec_last_iteration - s_vec_shocked, s_not_i_diff = s_not_i_vec_last_iteration - s_not_i_vec_shocked, dist_from_shocked=dist_from_shocked_person, iteration = "general")
  
  tmp<-rbind(initial_gain_df, final_gain_df)
}
  

